"""
Load balancer for MCP instances in CodeMCP gateway.

Manages multiple instances of MCP servers and distributes requests
across them for optimal performance and reliability.
"""

import asyncio
import time
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
from enum import Enum
import random
from datetime import datetime, timedelta
import aiohttp
import json

from ..core.error_handler import CodeMCPError, log_info, log_warning, log_debug, handle_error
from ..core.config import Config


class LoadBalancingStrategy(Enum):
    """Load balancing strategies."""
    ROUND_ROBIN = "round_robin"
    LEAST_CONNECTIONS = "least_connections"
    WEIGHTED_ROUND_ROBIN = "weighted_round_robin"
    RESPONSE_TIME = "response_time"
    RANDOM = "random"


class InstanceStatus(Enum):
    """MCP instance status."""
    HEALTHY = "healthy"
    UNHEALTHY = "unhealthy"
    DISABLED = "disabled"
    UNKNOWN = "unknown"


@dataclass
class MCPInstance:
    """Represents an MCP server instance."""
    id: str
    url: str
    language: str
    weight: float = 1.0
    max_connections: int = 100
    timeout_seconds: int = 30
    
    # Runtime state
    status: InstanceStatus = InstanceStatus.UNKNOWN
    current_connections: int = 0
    total_requests: int = 0
    failed_requests: int = 0
    avg_response_time: float = 0.0
    last_health_check: Optional[datetime] = None
    last_error: Optional[str] = None
    
    # Performance metrics
    response_times: List[float] = field(default_factory=list)
    error_count: int = 0
    consecutive_failures: int = 0
    
    def __post_init__(self):
        self.created_at = datetime.utcnow()
    
    @property
    def health_score(self) -> float:
        """Calculate health score (0.0-1.0)."""
        if self.total_requests == 0:
            return 1.0
        
        success_rate = 1.0 - (self.failed_requests / self.total_requests)
        connection_load = 1.0 - (self.current_connections / self.max_connections)
        
        # Factor in consecutive failures
        failure_penalty = min(0.1 * self.consecutive_failures, 0.8)
        
        return max(0.0, (success_rate * 0.6 + connection_load * 0.4) - failure_penalty)
    
    @property
    def is_available(self) -> bool:
        """Check if instance is available for requests."""
        return (self.status == InstanceStatus.HEALTHY and 
                self.current_connections < self.max_connections)
    
    def add_response_time(self, response_time: float):
        """Add a response time measurement."""
        self.response_times.append(response_time)
        
        # Keep only last 100 measurements
        if len(self.response_times) > 100:
            self.response_times = self.response_times[-100:]
        
        # Update average
        self.avg_response_time = sum(self.response_times) / len(self.response_times)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert instance to dictionary."""
        return {
            "id": self.id,
            "url": self.url,
            "language": self.language,
            "weight": self.weight,
            "status": self.status.value,
            "current_connections": self.current_connections,
            "total_requests": self.total_requests,
            "failed_requests": self.failed_requests,
            "avg_response_time": self.avg_response_time,
            "health_score": self.health_score,
            "last_health_check": self.last_health_check.isoformat() if self.last_health_check else None,
            "last_error": self.last_error,
            "consecutive_failures": self.consecutive_failures
        }


class LoadBalancer:
    """Load balancer for MCP instances."""
    
    def __init__(self, 
                 config: Config,
                 strategy: LoadBalancingStrategy = LoadBalancingStrategy.WEIGHTED_ROUND_ROBIN):
        self.config = config
        self.strategy = strategy
        self.instances: Dict[str, MCPInstance] = {}
        self.round_robin_index = 0
        
        # Health checking
        self.health_check_interval = 30  # seconds
        self.health_check_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        # Statistics
        self.stats = {
            "total_requests": 0,
            "failed_requests": 0,
            "avg_response_time": 0.0,
            "requests_per_second": 0.0,
            "last_reset": datetime.utcnow()
        }
        
        log_info(f"LoadBalancer initialized with strategy: {strategy.value}")
    
    async def start(self):
        """Start the load balancer."""
        if self.is_running:
            return
        
        self.is_running = True
        self.health_check_task = asyncio.create_task(self._health_check_loop())
        log_info("LoadBalancer started")
    
    async def stop(self):
        """Stop the load balancer."""
        if not self.is_running:
            return
        
        self.is_running = False
        
        if self.health_check_task:
            self.health_check_task.cancel()
            try:
                await self.health_check_task
            except asyncio.CancelledError:
                pass
        
        log_info("LoadBalancer stopped")
    
    def register_instance(self, instance: MCPInstance):
        """Register an MCP instance."""
        self.instances[instance.id] = instance
        log_info(f"Registered MCP instance: {instance.id} ({instance.language}) at {instance.url}")
    
    def unregister_instance(self, instance_id: str):
        """Unregister an MCP instance."""
        if instance_id in self.instances:
            instance = self.instances.pop(instance_id)
            log_info(f"Unregistered MCP instance: {instance_id}")
        else:
            log_warning(f"Attempted to unregister unknown instance: {instance_id}")
    
    def get_instance(self, instance_id: str) -> Optional[MCPInstance]:
        """Get instance by ID."""
        return self.instances.get(instance_id)
    
    def list_instances(self, 
                      language: Optional[str] = None, 
                      status: Optional[InstanceStatus] = None) -> List[MCPInstance]:
        """List instances with optional filters."""
        instances = list(self.instances.values())
        
        if language:
            instances = [i for i in instances if i.language == language]
        
        if status:
            instances = [i for i in instances if i.status == status]
        
        return instances
    
    async def select_instance(self, 
                            language: Optional[str] = None,
                            exclude_instances: List[str] = None) -> Optional[MCPInstance]:
        """Select an instance based on the load balancing strategy."""
        exclude_instances = exclude_instances or []
        
        # Filter available instances
        candidates = [
            instance for instance in self.instances.values()
            if (instance.is_available and 
                instance.id not in exclude_instances and
                (language is None or instance.language == language))
        ]
        
        if not candidates:
            log_warning(f"No available instances for language: {language}")
            return None
        
        # Apply load balancing strategy
        if self.strategy == LoadBalancingStrategy.ROUND_ROBIN:
            return self._select_round_robin(candidates)
        elif self.strategy == LoadBalancingStrategy.LEAST_CONNECTIONS:
            return self._select_least_connections(candidates)
        elif self.strategy == LoadBalancingStrategy.WEIGHTED_ROUND_ROBIN:
            return self._select_weighted_round_robin(candidates)
        elif self.strategy == LoadBalancingStrategy.RESPONSE_TIME:
            return self._select_by_response_time(candidates)
        elif self.strategy == LoadBalancingStrategy.RANDOM:
            return self._select_random(candidates)
        else:
            # Default to round robin
            return self._select_round_robin(candidates)
    
    def _select_round_robin(self, candidates: List[MCPInstance]) -> MCPInstance:
        """Round robin selection."""
        if not candidates:
            raise CodeMCPError("No candidates available for round robin selection")
        
        instance = candidates[self.round_robin_index % len(candidates)]
        self.round_robin_index += 1
        return instance
    
    def _select_least_connections(self, candidates: List[MCPInstance]) -> MCPInstance:
        """Select instance with least connections."""
        return min(candidates, key=lambda i: i.current_connections)
    
    def _select_weighted_round_robin(self, candidates: List[MCPInstance]) -> MCPInstance:
        """Weighted round robin based on health score and weight."""
        if not candidates:
            raise CodeMCPError("No candidates available for weighted selection")
        
        # Calculate weighted scores
        weighted_candidates = [
            (instance, instance.weight * instance.health_score)
            for instance in candidates
        ]
        
        # Sort by weighted score (descending)
        weighted_candidates.sort(key=lambda x: x[1], reverse=True)
        
        # Use weighted random selection
        total_weight = sum(weight for _, weight in weighted_candidates)
        if total_weight == 0:
            return random.choice(candidates)
        
        random_value = random.uniform(0, total_weight)
        cumulative_weight = 0
        
        for instance, weight in weighted_candidates:
            cumulative_weight += weight
            if random_value <= cumulative_weight:
                return instance
        
        # Fallback to first candidate
        return weighted_candidates[0][0]
    
    def _select_by_response_time(self, candidates: List[MCPInstance]) -> MCPInstance:
        """Select instance with best response time."""
        # Prioritize instances with response time data
        with_times = [i for i in candidates if i.response_times]
        without_times = [i for i in candidates if not i.response_times]
        
        if with_times:
            return min(with_times, key=lambda i: i.avg_response_time)
        else:
            return random.choice(without_times)
    
    def _select_random(self, candidates: List[MCPInstance]) -> MCPInstance:
        """Random selection."""
        return random.choice(candidates)
    
    async def execute_request(self, 
                            instance: MCPInstance,
                            request_data: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
        """Execute a request on the selected instance."""
        start_time = time.time()
        
        try:
            # Increment connection count
            instance.current_connections += 1
            
            # Make HTTP request to MCP instance
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=instance.timeout_seconds)) as session:
                async with session.post(
                    f"{instance.url}/execute",
                    json=request_data,
                    headers={"Content-Type": "application/json"}
                ) as response:
                    response_data = await response.json()
                    
                    # Calculate response time
                    response_time = time.time() - start_time
                    instance.add_response_time(response_time)
                    
                    # Update statistics
                    instance.total_requests += 1
                    self.stats["total_requests"] += 1
                    
                    if response.status == 200:
                        instance.consecutive_failures = 0
                        instance.last_error = None
                        return True, response_data
                    else:
                        # Handle HTTP error
                        instance.failed_requests += 1
                        instance.consecutive_failures += 1
                        instance.last_error = f"HTTP {response.status}: {response_data.get('error', 'Unknown error')}"
                        self.stats["failed_requests"] += 1
                        
                        return False, {
                            "error": f"Request failed with status {response.status}",
                            "details": response_data
                        }
        
        except asyncio.TimeoutError:
            # Handle timeout
            instance.failed_requests += 1
            instance.consecutive_failures += 1
            instance.last_error = "Request timeout"
            self.stats["failed_requests"] += 1
            
            error_response = handle_error(
                CodeMCPError("Request timeout"), 
                {"instance_id": instance.id, "timeout": instance.timeout_seconds}
            )
            return False, error_response
        
        except Exception as e:
            # Handle other errors
            instance.failed_requests += 1
            instance.consecutive_failures += 1
            instance.last_error = str(e)
            self.stats["failed_requests"] += 1
            
            error_response = handle_error(e, {"instance_id": instance.id})
            return False, error_response
        
        finally:
            # Decrement connection count
            instance.current_connections = max(0, instance.current_connections - 1)
            
            # Update average response time statistics
            if self.stats["total_requests"] > 0:
                total_time = sum(
                    sum(inst.response_times) for inst in self.instances.values()
                )
                total_measurements = sum(
                    len(inst.response_times) for inst in self.instances.values()
                )
                if total_measurements > 0:
                    self.stats["avg_response_time"] = total_time / total_measurements
    
    async def _health_check_loop(self):
        """Periodic health check loop."""
        while self.is_running:
            try:
                await self._perform_health_checks()
                await asyncio.sleep(self.health_check_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                log_warning(f"Error in health check loop: {e}")
                await asyncio.sleep(self.health_check_interval)
    
    async def _perform_health_checks(self):
        """Perform health checks on all instances."""
        health_check_tasks = [
            self._check_instance_health(instance)
            for instance in self.instances.values()
        ]
        
        if health_check_tasks:
            await asyncio.gather(*health_check_tasks, return_exceptions=True)
    
    async def _check_instance_health(self, instance: MCPInstance):
        """Check health of a single instance."""
        try:
            start_time = time.time()
            
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) as session:
                async with session.get(f"{instance.url}/health") as response:
                    response_time = time.time() - start_time
                    
                    if response.status == 200:
                        health_data = await response.json()
                        
                        # Update instance status
                        if health_data.get("status") == "healthy":
                            instance.status = InstanceStatus.HEALTHY
                            instance.consecutive_failures = 0
                        else:
                            instance.status = InstanceStatus.UNHEALTHY
                    else:
                        instance.status = InstanceStatus.UNHEALTHY
                        instance.consecutive_failures += 1
                    
                    instance.last_health_check = datetime.utcnow()
                    instance.add_response_time(response_time)
        
        except Exception as e:
            instance.status = InstanceStatus.UNHEALTHY
            instance.consecutive_failures += 1
            instance.last_error = f"Health check failed: {str(e)}"
            instance.last_health_check = datetime.utcnow()
            
            log_debug(f"Health check failed for {instance.id}: {e}")
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get load balancer statistics."""
        now = datetime.utcnow()
        time_elapsed = (now - self.stats["last_reset"]).total_seconds()
        
        if time_elapsed > 0:
            self.stats["requests_per_second"] = self.stats["total_requests"] / time_elapsed
        
        return {
            "strategy": self.strategy.value,
            "total_instances": len(self.instances),
            "healthy_instances": len([i for i in self.instances.values() 
                                    if i.status == InstanceStatus.HEALTHY]),
            "total_requests": self.stats["total_requests"],
            "failed_requests": self.stats["failed_requests"],
            "success_rate": (1.0 - (self.stats["failed_requests"] / max(1, self.stats["total_requests"]))) * 100,
            "avg_response_time": self.stats["avg_response_time"],
            "requests_per_second": self.stats["requests_per_second"],
            "instances": [instance.to_dict() for instance in self.instances.values()]
        }
    
    def reset_statistics(self):
        """Reset load balancer statistics."""
        self.stats = {
            "total_requests": 0,
            "failed_requests": 0,
            "avg_response_time": 0.0,
            "requests_per_second": 0.0,
            "last_reset": datetime.utcnow()
        }
        
        # Reset instance statistics
        for instance in self.instances.values():
            instance.total_requests = 0
            instance.failed_requests = 0
            instance.response_times = []
            instance.avg_response_time = 0.0
            instance.error_count = 0
            instance.consecutive_failures = 0
        
        log_info("Load balancer statistics reset")
    
    async def handle_request(self, 
                           request_data: Dict[str, Any],
                           language: Optional[str] = None,
                           max_retries: int = 2) -> Dict[str, Any]:
        """Handle a request with load balancing and retries."""
        excluded_instances = []
        
        for attempt in range(max_retries + 1):
            # Select instance
            instance = await self.select_instance(language, excluded_instances)
            
            if not instance:
                return {
                    "success": False,
                    "error": "No available instances",
                    "attempt": attempt + 1
                }
            
            # Execute request
            success, response = await self.execute_request(instance, request_data)
            
            if success:
                return {
                    "success": True,
                    "result": response,
                    "instance_id": instance.id,
                    "attempt": attempt + 1
                }
            else:
                # Add instance to exclusion list for retries
                excluded_instances.append(instance.id)
                log_warning(f"Request failed on instance {instance.id}, attempt {attempt + 1}")
        
        # All attempts failed
        return {
            "success": False,
            "error": "All retry attempts failed",
            "attempts": max_retries + 1,
            "last_response": response if 'response' in locals() else None
        }